import CodeSnapped from "./treesitter.ts";
import { CodeSnippetsCodebaseIndex } from "./core/codeSnippetsIndex.js";
import { saveLancedb, searchLancedb } from "./core/embed.js";

/**
 * 整合流程：解析代码、存储到SQLite和LanceDB
 */
async function processCodebase() {
  try {
    const fileDir = "./src";
    const CodeSnappeder = new CodeSnapped();
    
    console.log("开始解析代码文件...");
    // 解析代码文件获取代码块
    const codeChunks = await CodeSnappeder.chunkCode(fileDir);
    console.log(`成功解析到 ${codeChunks.length} 个代码块`);
    
    if (codeChunks.length === 0) {
      console.log("没有找到代码块，结束处理");
      return;
    }
    
    // 保存到SQLite数据库
    console.log("正在保存代码片段到SQLite数据库...");
    await CodeSnippetsCodebaseIndex.add(codeChunks);
    console.log("代码片段保存成功");
    
    // 准备向量数据库数据格式
    const vectorData = codeChunks.map(chunk => ({
      text: chunk.code || '', // 使用代码内容作为嵌入文本
      path: chunk.filePath || '',
      signature: chunk.signature || '',
      startLine: chunk.startPosition || 0,
      endLine: chunk.endPosition || 0
    }));
    
    // 保存到LanceDB向量数据库
    console.log("正在生成向量并保存到LanceDB...");
    await saveLancedb(vectorData);
    console.log("向量数据保存成功");
    
    // 示例查询
    await exampleQuery();
    
  } catch (error) {
    console.error("处理代码库时出错:", error);
  }
}

/**
 * 示例查询函数
 */
async function exampleQuery() {
  try {
    // 示例1: 使用SQLite查询特定签名的代码
    const signature = "getDB";
    console.log(`\n查询签名为 "${signature}" 的代码片段:`);
    const sqliteResults = await CodeSnippetsCodebaseIndex.getAll(signature);
    console.log(sqliteResults);
    
    // 示例2: 使用LanceDB进行语义搜索
    const queryText = "创建数据库连接";
    const topK = 5;
    console.log(`\n语义搜索 "${queryText}" 的前 ${topK} 个结果:`);
    const vectorResults = await searchLancedb(queryText, topK);
    vectorResults.forEach((result, index) => {
      console.log(`\n结果 ${index + 1}:`);
      // 限制输出长度以便查看
      const displayText = result.length > 200 ? result.substring(0, 200) + '...' : result;
      console.log(displayText);
    });
    
  } catch (error) {
    console.error("查询时出错:", error);
  }
}

// 执行主函数
processCodebase();